% Copyright (C) 2019-2023 Benjamin Born, Francesco D'Ascanio, Gernot J. Mueller, Johannes Pfeifer
%
% This is free software: you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published by
% the Free Software Foundation, either version 3 of the License, or
% (at your option) any later version.
%
% It is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
% 
% For a copy of the GNU General Public License,
% see <http://www.gnu.org/licenses/>.
 
 
%Compute impulse responses to positive and negative g shocks for the float
%case

clear; clc;
if ~isfolder('Figures')
    mkdir('.','Figures');
end
if ~isfolder('IRF_results')
    mkdir('.','IRF_results');
end

addpath('Auxiliary_functions/')

N_periods = 150; %number of simulation periods (quarters)
N_rep = 1e6; %number of replications
T_ergodic = 1e6; %periods for simulation of mean;
T_burnin = 1e+3; %burn-in period

%load parameter values
loaded_PFI=load('Policy_functions/vfi_float_g_low_beta_large_grid.mat');
omega=loaded_PFI.par.omega;
alfa=loaded_PFI.par.alfa;
hbar=loaded_PFI.par.hbar;
xi=loaded_PFI.par.xi;


%% Set initial conditions of exogenous states, debt, and past wage equal to their respective unconditional means
temp=unique(loaded_PFI.grid.y_level);
y_0 = temp(round(length(temp)/2));
temp=unique(loaded_PFI.grid.r_level);
r_0 = temp(round(length(temp)/2));
temp=unique(loaded_PFI.grid.g_level);
g_0 = temp(round(length(temp)/2));
g_0_pos = loaded_PFI.grid.g_level(find(loaded_PFI.grid.g_level==g_0,1,'last')+1);
g_0_neg = loaded_PFI.grid.g_level(find(loaded_PFI.grid.g_level==g_0, 1 )-1);
  
d_index = 362;
d_0 = loaded_PFI.grid.d_grid(d_index);


%store initial values
y_T_start = repmat(y_0,[1 3]);
r_start = repmat(r_0,[1 3]);
g_N_start = [g_0 g_0_pos g_0_neg];
d_start = repmat(d_0,[1 3]);

c_T = y_0 + d_0./(1+r_0) - d_0;
w_lag_start=(1-omega)/omega*((hbar^alfa - g_0)./c_T).^(-1/xi) *alfa*hbar^(alfa-1);%full-employment real wage rate in terms of tradables
y_N = hbar.^alfa; %nontradable output
c_N = y_0 - g_0; %nontradable consumption
p_lag_start = (1-omega)/omega*(c_N./c_T).^(-1/xi); %relative price of nontradables in terms of tradables
p_lag_start = repmat(p_lag_start,[1 3]);
w_lag_start = repmat(w_lag_start,[1 3]);
clear('c_T','y_N','c_N','r_0','g_0','d_0','g_0_pos','g_0_neg');

%P_trans is the transition probability matrix of tradable output. Cumul_P_trans is the cumulative probability matrix (useful for drawing tradable output realizations)
Cum_prob = cumsum(loaded_PFI.MC.P_trans,2);

[Y_base, var_name, pos]=simulate_model_float(y_T_start(1),r_start(1),g_N_start(1),d_start(1),w_lag_start(1),p_lag_start(1),T_ergodic,loaded_PFI.dp,loaded_PFI.grid,loaded_PFI.par,Cum_prob,1);
Y_base_mean=mean(Y_base(T_burnin+1:end,:));

nvar = length(fieldnames(pos)); %number of variables for which IR are computed

%Initialize vector of impulse responses
Y = zeros(N_periods,nvar,N_rep);
Y_pos = zeros(N_periods,nvar,N_rep);
Y_neg = zeros(N_periods,nvar,N_rep);

%reset random number generator
rng('default')
tic
for replic_iter=1:N_rep
    %set initial condition for current replication
    y_T = y_T_start;
    r = r_start;
    g_N = g_N_start;
    d = d_start;
    w_lag= w_lag_start;
    p_lag= p_lag_start;
    
    Y(:,:,replic_iter)=simulate_model_float(y_T(1),r(1),g_N(1),d(1),w_lag(1),p_lag(1),N_periods,loaded_PFI.dp,loaded_PFI.grid,loaded_PFI.par,Cum_prob);
    Y_pos(:,:,replic_iter)=simulate_model_float(y_T(2),r(2),g_N(2),d(2),w_lag(2),p_lag(2),N_periods,loaded_PFI.dp,loaded_PFI.grid,loaded_PFI.par,Cum_prob);
    Y_neg(:,:,replic_iter)=simulate_model_float(y_T(3),r(3),g_N(3),d(3),w_lag(3),p_lag(3),N_periods,loaded_PFI.dp,loaded_PFI.grid,loaded_PFI.par,Cum_prob);
    
end
toc

IRFs=NaN(nvar,N_periods,2);
%% Plot IRFs
h1=figure('Name','Positive shock');
for var_iter=1:min(nvar,16)
    subplot(4,4,var_iter)
    IRFs(var_iter,:,1)=mean(squeeze((Y_pos(:,var_iter,:)-Y(:,var_iter,:))),2);
    plot(IRFs(var_iter,:,1)./Y_base_mean(var_iter))
    xlim([1 N_periods]);
    title(var_name{var_iter}, 'interpreter','latex')
end

print('IRF_results/IRF_float_positive','-depsc2')
saveas(h1,'IRF_results/IRF_float_positive')

if nvar>16
    h3=figure('Name','Positive shock 2');
    for var_iter=1:min(nvar-16,16)
        subplot(4,4,var_iter)
        IRFs(16+var_iter,:,1)=mean(squeeze((Y_pos(:,16+var_iter,:)-Y(:,16+var_iter,:))),2);
        plot(IRFs(16+var_iter,:,1)./Y_base_mean(16+var_iter))
        xlim([1 N_periods]);
        title(var_name{16+var_iter}, 'interpreter','latex')
    end
end
saveas(h3,'IRF_results/IRF_float_positive_2')

h2=figure('Name','Negative shock');
for var_iter=1:min(nvar,16)
    subplot(4,4,var_iter)
    IRFs(var_iter,:,2)=mean(squeeze((Y_neg(:,var_iter,:)-Y(:,var_iter,:))),2);
    plot(IRFs(var_iter,:,2)./Y_base_mean(var_iter))
    xlim([1 N_periods]);
    title(var_name{var_iter}, 'interpreter','latex')
end
saveas(h2,'IRF_results/IRF_float_negative')

if nvar>16
    h4=figure('Name','Negative shock 2');
    for var_iter=1:min(nvar-16,16)
        subplot(4,4,var_iter)
        IRFs(16+var_iter,:,2)=mean(squeeze((Y_neg(:,16+var_iter,:)-Y(:,16+var_iter,:))),2);
        plot(IRFs(16+var_iter,:,2)./Y_base_mean(16+var_iter))
        xlim([1 N_periods]);
        title(var_name{16+var_iter}, 'interpreter','latex')
    end    
end
saveas(h4,'IRF_results/IRF_float_negative_2')

save('IRF_results/irfs_float.mat','IRFs','Y_base_mean','pos','var_name','-v7.3')